/**
 * @file discrete_distr.h
 *
 * Discrete distributions
 * 
 * @author Dahua Lin
 */

#ifndef SLIPP_DISCRETE_DISTR_H_
#define SLIPP_DISCRETE_DISTR_H_

#include <slipp/prob_base/pdistr_base.h>

#include <limits>

namespace sli
{

	template<class RStream, typename IntType>
	class discrete_sampler
	{
	public:
		SLIPP_DEFINE_SCALAR_DISTR_TYPES(IntType)

		discrete_sampler(RStream& rstream, const IntType& n, const double *w, double tw)
		: m_rng(rstream, n, w, tw)
		{
		}

		sample_type draw()
		{
			return m_rng.get();
		}

		sample_collection_type draw(size_t n)
		{
			sample_collection_type x((index_t)n);
			m_rng.get_m(n, x.data());
			return x;
		}

	private:
		discrete_rng<RStream, IntType> m_rng;
	};


	template<class RStream, typename IntType>
	class sorted_discrete_sampler
	{
	public:
		SLIPP_DEFINE_SCALAR_DISTR_TYPES(IntType)

		sorted_discrete_sampler(RStream& rstream, const IntType& n, const double *w, double tw)
		: m_sorted_weights((index_t)n)
		, m_sorted_ints((index_t)n)
		, m_rng(rstream, (int32_t)n, m_sorted_weights.data(), tw)
		{
			Eigen::Map<darow_t> mw(const_cast<double*>(w), (index_t)n);
			sort_with_index_des(mw, m_sorted_weights.get(), m_sorted_ints.get());
		}

		const darow_t& sorted_weights() const
		{
			return m_sorted_weights.get();
		}

		const EIGEN_AROW(IntType)& sorted_values() const
		{
			return m_sorted_ints.get();
		}

		sample_type draw()
		{
			int32_t i = m_rng.get();
			return m_sorted_ints.get().coeff(i);
		}

		sample_collection_type draw(size_t n)
		{
			iarow_t I((index_t)n);
			m_rng.get_m(n, I.data());
			return select_coeffs(m_sorted_ints.get(), I);
		}

	private:
		share_darow_t m_sorted_weights;
		SLI_SAROW(IntType) m_sorted_ints;

		discrete_rng<RStream, int32_t> m_rng;
	};


	/**
	 * The discrete distribution
	 */
	template<typename IntType=int32_t>
	class discrete_distr
	{
	public:
		SLIPP_DEFINE_SCALAR_DISTR_TYPES(IntType)

		class param_type
		{
		public:
			param_type(share_darow_t pvec)
			: m_pvec(pvec)
			{
			}

			IntType K() const
			{
				return (IntType)m_pvec.size();
			}

			const darow_t& pvec() const
			{
				return m_pvec.get();
			}

			double p(IntType x) const
			{
				return m_pvec.get().coeff(x);
			}

		private:
			share_darow_t m_pvec;
		};

	public:
		// constructor

		discrete_distr(share_darow_t pvec)
		: m_param(pvec)
		{
		}

		discrete_distr(const param_type& param)
		: m_param(param)
		{
		}

		// parameters

		const param_type& param() const
		{
			return m_param;
		}

		IntType K() const
		{
			return m_param.K();
		}

		const darow_t& pvec() const
		{
			return m_param.pvec();
		}

		double p(IntType x) const
		{
			return m_param.p(x);
		}

		// statistics

		double mean() const
		{
			IntType k = K();
			const double *pv = pvec().data();

			double v = 0;
			for (IntType i = 1; i < k; ++i)
			{
				v += pv[i] * i;
			}
			return v;
		}

		double entropy() const
		{
			IntType k = K();
			const double *pv = pvec().data();

			double e = 0;
			for (IntType i = 0; i < k; ++i)
			{
				if (pv[i] > 0) e += pv[i] * math::log(pv[i]);
			}

			return -e;
		}

	public:

		class pdf_evaluator_type
		{
		public:
			typedef discrete_distr::sample_type sample_type;
			typedef discrete_distr::sample_collection_type sample_collection_type;

			pdf_evaluator_type(const param_type& param)
			: m_pvec(param.pvec()), m_K(param.K())
			{
			}

			SLIPP_ENSURE_INLINE double pdf(const sample_type& x) const
			{
				return in_range(x) ? _p(x) : 0.0;
			}

			SLIPP_ENSURE_INLINE double logpdf(const sample_type& x) const
			{
				return in_range(x) ? math::log(_p(x)) : -std::numeric_limits<double>::infinity();
			}

			darow_t pdf(const sample_collection_type& X) const
			{
				return map(_detail::make_pdf_eval_functor(*this), X);
			}

			darow_t logpdf(const sample_collection_type& X) const
			{
				return map(_detail::make_logpdf_eval_functor(*this), X);
			}

		private:
			SLIPP_ENSURE_INLINE bool in_range(const sample_type& x) const
			{
				return x >= 0 && x < m_K;
			}

			SLIPP_ENSURE_INLINE double _p(const sample_type& x) const
			{
				return m_pvec.coeff((index_t)x);
			}

		private:
			const darow_t& m_pvec;
			sample_type m_K;
		};

		pdf_evaluator_type get_pdf_evaluator() const
		{
			return pdf_evaluator_type(m_param);
		}

	public:
		// sampling

		template<class RStream>
		discrete_sampler<RStream, IntType> get_sampler(RStream& rstream) const
		{
			return discrete_sampler<RStream, IntType>(rstream, m_param.K(), m_param.pvec().data(), 1.0);
		}

		template<class RStream>
		sorted_discrete_sampler<RStream, IntType> get_sorted_sampler(RStream& rstream) const
		{
			return sorted_discrete_sampler<RStream, IntType>(rstream, m_param.K(), m_param.pvec().data(), 1.0);
		}

	private:
		param_type m_param;

	}; // end class discrete_distr


	template<typename IntType> struct is_pdistr<discrete_distr<IntType> >
	{
		static const bool value = true;
	};

	template<typename IntType>
	struct pdistr_traits<discrete_distr<IntType> >
	: public pdistr_traits_base<discrete_distr<IntType>, true, true> { };


	template<class RStream, typename IntType>
	struct sampler_of<discrete_distr<IntType>, RStream>
	{
		typedef discrete_sampler<RStream, IntType> type;
	};

	template<typename IntType>
	struct pdf_evaluator_of<discrete_distr<IntType> >
	{
		typedef typename discrete_distr<IntType>::pdf_evaluator_type type;
	};

}

#endif


